#PBAR_CreateStandardSizeImages.py
#
#   Make standard cargo and zspec size images, modified markers that work with these new images.
#
# 

import PBAR_Zspec, PBAR_Cargo
import numpy as np
import os
import matplotlib.pyplot as plt
from scipy import interpolate
from matplotlib import cm
from scipy import ndimage
import copy
import glob
import shutil
import codecs

setNum = 2

# set 1
#dataDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScans'
#dataOutputBaseDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScansStandardWidth'

# set 2
dataDir = r'E:\PBAR\data4\BasicScans'
dataOutputBaseDir = r'E:\PBAR\data4\BasicScansStandardWidth'

# define parameters for standard width images
numberTimeSlices = 2000
preCargoTimeSlices = 70

#################################
## DEFINE FILE NAMES

cargoConfigBaseDirList = glob.glob(os.path.join(dataDir, '*'))
#cargoConfigBaseDirList = cargoConfigBaseDirList[0:5]

for (dirIndex, cargoConfigBaseDir) in enumerate(cargoConfigBaseDirList):
    subDirList = glob.glob(os.path.join(cargoConfigBaseDir, '*'))
    for (subDirIndex, dataPath) in enumerate(subDirList):
        print(dataPath)

        a, b = os.path.split(dataPath)
        a, c = os.path.split(a)

        dataOutputDir = os.path.join(dataOutputBaseDir, c, b)
        
        # make paths if not present
        if not os.path.exists(os.path.join(dataOutputBaseDir, c)):
            os.mkdir(os.path.join(dataOutputBaseDir, c))

        if not os.path.exists(dataOutputDir):
            os.mkdir(dataOutputDir)
        
        zspecScanNumberList = []
        cargoScanNumberList = []
        
        fullFilenameList = []
        filenameList = []
        fullFilenameOutputList = []
        
        filenameCargoList = []
        fullFilenameCargoList = []
        fullFilenameCargoOutputList = []
        markerfullFilenameCargoList = []
        markerfullFilenameCargoOutputList = []        
        
        # find the zspec scan number
        temp = glob.glob(os.path.join(dataPath, '*.npy'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No npy file found.  Skipping.')
            continue
        a, filename = os.path.split(temp[0])
        zspecScanNumberList.append(filename[0:4])
        filenameList.append(filename)
        fullFilenameList.append(os.path.join(dataPath, filename))
        fullFilenameOutputList.append(os.path.join(dataOutputDir, filename))
        

        # get the cargo image name
        temp = glob.glob(os.path.join(dataPath, '*.cargoimage'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No cargoimage file found.  Skipping.')
            continue
        a, filename = os.path.split(temp[0])
        cargoScanNumberList.append(filename[0:16])
        filenameCargoList.append(filename)    
        fullFilenameCargoList.append(os.path.join(dataPath, filename))
        
        fullFilenameCargoOutputList.append(os.path.join(dataOutputDir, filename))
        
        markerfullFilenameCargoList.append(os.path.join(dataPath, filename.replace('cargoimage', 'cargomarker')))

        markerfullFilenameCargoOutputList.append(os.path.join(dataOutputDir, filename.replace('cargoimage', 'cargomarker')))
        
        ######################
        ## READ IN DATA
        
        # Read in marker files
        markerList = []
        for (index, f) in enumerate(markerfullFilenameCargoList):
            markerList.append(PBAR_Cargo.ReadCargoMarker(f))
        
        # Read Zspec
        energyList = []
        dat = []
        datStandardWidth = []
        startBinZspecList = []
        for (index, fullFilename) in enumerate(fullFilenameList):
            (a, b) = PBAR_Zspec.ReadZspecBasicScanNumpy(fullFilename)
            energyList.append(a)
            dat.append(b)


            if setNum == 1:
                # for set 1        
                startBinZspec = PBAR_Zspec.FindBasicZspecStart(b)
            elif setNum == 2:
                # for set 2
                startBinZspec = PBAR_Zspec.FindBasicZspecStartVer2(b)
            
            startBinZspecList.append(startBinZspec)
            
            temp = b[(startBinZspec-preCargoTimeSlices):,:][0:numberTimeSlices,:,:] # this could have length less than numberTimeSlices    
            
           # This ensure that we have standard size images    
            temp2 = np.zeros((numberTimeSlices, b.shape[1], b.shape[2]))
            if temp.shape[0] < numberTimeSlices:
                temp2[0:temp.shape[0],:,:] = temp
            else:
                temp2 = temp
            datStandardWidth.append(temp2)
            # write standard width file
            temp2 = temp2.astype(np.uint16)
            print('Wrote: %s' %fullFilenameOutputList[index])
            np.save(fullFilenameOutputList[index], temp2)
        del a, b
        
        # Read Cargo
        numberZspecDetectors = dat[0].shape[1]
        datCargo = []
        datCargoStandardWidth = []
        startBinCargoList = []
        
        for (index, f) in enumerate(fullFilenameCargoList):
            (A,bpp,formatt,flag,low1,high1,low2,high2) = PBAR_Cargo.ReadCargoImage(f)
            datCargo.append(A)
        
            # this should do the same thing
            indices1 = 1 + 4 * np.arange(136) 
            indices2 = 2 + 4* np.arange(136)
            temp = (A[:,indices1] + A[:,indices2])/2
            
            startBinCargo = PBAR_Cargo.FindCargoStart(A)
            startBinCargoList.append(startBinCargo)
            
            # cut section in front and limit to numberTimeSlices number of time slices
            temp = temp[(startBinCargo-preCargoTimeSlices):,:][0:numberTimeSlices,:]
            
            # This ensure that we have standard size images
            # excess is filled with zeros
            temp2 = np.zeros((numberTimeSlices, numberZspecDetectors))
            if temp.shape[0] < numberTimeSlices:
                temp2[0:temp.shape[0],:] = temp
            else:
                temp2 = temp[0:numberTimeSlices,:]
                
            datCargoStandardWidth.append(temp2)
            
            # write standard width file
            print('Wrote %s' %fullFilenameCargoOutputList[index])
            np.save(fullFilenameCargoOutputList[index], temp2)
            
            # write to file
            
        del A,bpp,formatt,flag,low1,high1,low2,high2
        del indices1, indices2, temp, temp2
        
        # Modify the marker files based on found start times
        markerStandardWidthList = []
        
        for (index, markers) in enumerate(markerList):
            # first copy over
            temp = copy.deepcopy(markers)
            
            # adjust the x values
            offset = (startBinCargo - preCargoTimeSlices)
            for i in np.arange(len(temp)):
                temp[i]['rec_left'] = temp[i]['rec_left'] - offset
                temp[i]['rec_right'] = temp[i]['rec_right'] - offset
                temp[i]['x'] = temp[i]['x'] - offset# write standard width file
                # see if 'left and 'right exist
                if 'left' in temp[i]:
                    temp[i]['left']['x'] = temp[i]['left']['x'] - offset
                if 'right' in temp[i]:
                    temp[i]['right']['x'] = temp[i]['right']['x'] - offset
            # adjust the y values
            for i in np.arange(len(temp)):
                temp[i]['rec_top'] = np.round((temp[i]['rec_top'] - 1.5) / 4.0)
                temp[i]['rec_bottom'] = np.round((temp[i]['rec_bottom'] - 1.5) / 4.0)
                temp[i]['y'] = np.round((temp[i]['y'] - 1.5) / 4.0)
                # see if 'left and 'right exist
                if 'left' in temp[i]:
                    temp[i]['left']['y'] = np.round((temp[i]['left']['y'] - 1.5) / 4.0)
                if 'right' in temp[i]:
                    temp[i]['right']['y'] = np.round((temp[i]['right']['y'] - 1.5) / 4.0)
            markerStandardWidthList.append(copy.copy(temp))
        

        # save modified marker file
        # read cargomarker, adust offset, save to new marker file
        
        filename = markerfullFilenameCargoList[0]
        offset = (startBinCargo-preCargoTimeSlices)

        filenameNew = markerfullFilenameCargoOutputList[0]  
        
        # open old file for reading
        with codecs.open(filename, 'rb', encoding = 'utf-16') as fidOld:
            # open file for writting
            with codecs.open(filenameNew, 'wb', encoding = 'utf-16') as fid:
                # Check to see if the file has been convereted already by checking for existence of cargomarkerOLD file.
                for lineIn in fidOld:
                    if len(lineIn) < 3:
                        fid.write(lineIn) # the first line seems to contain some whitespace
                    else:
                        lineIn2 = copy.copy(lineIn)
                        # Modify the entires
                        preStringList = ['left=', 'right=', 'center x=', 'left x=', 'right x=']                    
                        # cycle through the list of x direction values.
                        for (index, preString) in enumerate(preStringList):
                            # check if it actually exists
                            if lineIn2.find(preString) > -1:
                                index1 = lineIn2.find(preString) + len(preString) + 1 # index of largest digit
                                index2 = lineIn2[index1:].find('"')+index1 # index parenthesis to the right of the smallest digit
                                theNumber = int(lineIn2[index1:index2]) - offset
                                lineIn2 = u'%s%d%s' %(lineIn2[:index1], theNumber, lineIn2[index2:])
                                
                        preStringList = ['top=', 'bottom=', ' y=']                    
                        # cycle through the list of x direction values.
                        for (index, preString) in enumerate(preStringList):
                            # check if it actually exists
                            if lineIn2.find(preString) > -1:
                                index1 = lineIn2.find(preString) + len(preString) + 1 # index of largest digit
                                index2 = lineIn2[index1:].find('"')+index1 # index parenthesis to the right of the smallest digit
                                theNumber = round((int(lineIn2[index1:index2]) - 1.5 ) /4.0 )
                                lineIn2 = u'%s%d%s' %(lineIn2[:index1], theNumber, lineIn2[index2:])
                                if lineIn2.find(preString, index1+1) > -1:
                                    index1 = lineIn2.find(preString, index1+1) + len(preString) + 1 # index of largest digit
                                    index2 = lineIn2[index1:].find('"')+index1 # index parenthesis to the right of the smallest digit
                                    theNumber = round((int(lineIn2[index1:index2]) - 1.5 ) /4.0 )
                                    lineIn2 = u'%s%d%s' %(lineIn2[:index1], theNumber, lineIn2[index2:])
                                    if lineIn2.find(preString, index1+1) > -1:
                                        index1 = lineIn2.find(preString, index1+1) + len(preString) + 1 # index of largest digit
                                        index2 = lineIn2[index1:].find('"')+index1 # index parenthesis to the right of the smallest digit
                                        theNumber = round((int(lineIn2[index1:index2]) - 1.5 ) /4.0 )
                                        lineIn2 = u'%s%d%s' %(lineIn2[:index1], theNumber, lineIn2[index2:])
                                
                        #write the modified line to the new file
                        fid.write(lineIn2)
        fid.close()
        fidOld.close()
        print('Created: %s' %filename)
